Skip to content

Dynamic shape on the C++ trace path (elementwise add)#269

Open
YWHyuk wants to merge 46 commits into
feature/togsim-cpp-tracefrom
feature/togsim-dynamic-shape
Open

Dynamic shape on the C++ trace path (elementwise add)#269
YWHyuk wants to merge 46 commits into
feature/togsim-cpp-tracefrom
feature/togsim-dynamic-shape

Conversation

@YWHyuk

@YWHyuk YWHyuk commented Jun 23, 2026

Copy link
Copy Markdown
Collaborator

Summary

Adds torch.compile(dynamic=True) support on the C++ trace path. One compiled
trace producer .so serves any input size: the producer reads its loop bound
from shape_args at runtime, while the per-tile cost table (sampled once, since a
tile's cost is shape-invariant) keys the timing.

Verified end to end: a single dynamic a + b runs at N=1024 (2 tiles, 183
cycles) and N=2048 (4 tiles, 261 cycles) from the same .so, the trace
scaling with shape_args.

What each commit does

  1. Guard MLIR tile sizing against symbolic dims — neutralize the tile-fit
    heuristics for a sympy dim (they only minimize a known dim's tail); all gated
    on isinstance(.., sympy.Expr), so static is unchanged.
  2. Emit symbolic loop bounds and dynamic memref dims — size symbol becomes a
    scalar-int kernel arg; memref<?xf32>; affine.for ... to %<name>_bound with a
    top-of-function memref.load + index_cast prologue.
  3. Make the kernel meta import-safe — stringify sympy in arg_attributes so the
    generated wrapper imports (the extent already arrives as a runtime arg).
  4. Skip compile-time Spike validation for dynamic — its fixed-shape validation
    binary can't be instantiated for a runtime extent.
  5. Sample per-tile cycles on a one-tile copypin_loops_to_one_tile (general
    for static + dynamic) runs gem5 sampling on a one-tile copy; the symbolic IR is
    kept for the producer / cost table.
  6. Emit a dynamic-shape trace producerbuild_skeleton skips the loop_end
    serialization it never needed; lower_to_emitc re-sources each still-used size
    arg from shape_args[k] via emitc.subscript, so the producer loop reads
    for (iv=0; iv<shape_args[k]; iv+=step).
  7. Pass the runtime shape via the attribute file — the scalar size goes into the
    existing per-kernel attribute YAML (shape_args), run_standalone passes
    --attribute, and main.cc fills shape_args for run_producer.
  8. Testtests/ops/elementwise/test_dynamic_add.py.

Known limitations / follow-ups

  • Timing only: the Spike functional run (which fills the output tensor) is
    skipped for dynamic, so output values are zeros today; the test checks output
    shape. Functional output for dynamic is a follow-up.
  • Static cost sampling can use the (already general) pin_loops_to_one_tile once it
    is decoupled from run_tog; op coverage beyond contiguous 1D add (matmul /

    1D-dynamic: symbolic-stride addressing, multi-symbol shape_args) is future work.

  • One change to the loop-padding mlir-opt pass (skip a symbolic bound) lives in the
    LLVM fork and is not part of this PR.

Base: stacked on feature/togsim-cpp-trace (#267).

🤖 Generated with Claude Code

@YWHyuk YWHyuk force-pushed the feature/togsim-dynamic-shape branch from c2a243b to 6b39c94 Compare June 23, 2026 06:00
@YWHyuk YWHyuk force-pushed the feature/togsim-cpp-trace branch 2 times, most recently from b7c1ec4 to 9033945 Compare June 24, 2026 13:37
@YWHyuk YWHyuk force-pushed the feature/togsim-dynamic-shape branch from 8e9db1a to 82e9255 Compare June 24, 2026 14:43
@YWHyuk YWHyuk force-pushed the feature/togsim-cpp-trace branch from c166abd to ed5c747 Compare June 25, 2026 07:29
YWHyuk and others added 25 commits July 7, 2026 16:09
… feed

Skeleton + EmitC + cost/dep analysis on the frontend; the trace runtime,
loader, bridge, and Core feed on the simulator; shared MLIR pass helpers and
the pipeline tests.
Per-record tag key in the bridge plus per-iteration tag alloc in
dma-fine-grained so multi-tile-K and conv loads do not collide; strip the
reduction accum marker from the memory_barrier slot.
togsim_dispatch with TILE_BEGIN/TILE_END; outline each work-item into
togsim_kernel_tile.
DMA-capacity throttle and frozen-state guard, per-core VMEM in the configs,
and the SA weight-buffer throttle.
trace_timeline.py with per-work-item grouping and resource-centric DMA lanes;
the trace logs the first DRAM response and the assigned systolic array, and
scopes the compute barrier to its dispatch.
Default to the trace path; fix uninitialized Instruction fields, the matmul
accumulator wedge, fused-subtile dedup, nested/fused epilogue dataflow, and
dma_wait fusion; bound concurrent dispatches to the spad, round-robin
work-items within a partition, benchmark autotune and run the multi-tenant
scheduler through the trace path, and emit trace.so for pooling/reduction.
Carry simulator headers through the wrapper for cache-safe replay; drop verbose
[P3-trace] logs; fix the key.mlir compile race in load().
… runtime model

Replace the trace bridge's accumulated special cases with one dataflow rule and
clean up the runtime that consumes it.

Dependency rule: per SRAM buffer keep a writers SET; a reader depends on all
current writers (occupancy=ISSUE when both are systolic-array ops, else
latency=DONE); a writer REPLACEs the set. The only exception is is_mm_accum (a
matmul that reads and writes the same buffer = a commutative accumulator): skip
its read edge and UNION its write, waiting only the non-matmul init seed and not
ordering co-matmuls. This drops the matmul-accumulator chain that deadlocked the
SA weight-slot pipeline while keeping the init->matmul edge, and lets a vector
epilogue or the store wait every K matmul (fixes the pure-vector store that an
empty COMPUTE_BAR let slip).

Remove COMPUTE_BAR entirely: a matmul is its own DONE-handle (finish == SA
drain), so the store JOINs the matmul writers directly. The whole emit/loader
chain is gone -- build_skeleton, lower_to_emitc, togsim.compute_barrier, the
runtime symbol, the Opcode/case/_fence_finish, and TraceRec::COMPUTE_BAR -- so a
stale producer fails loudly instead of emitting records the bridge would drop.
Only MEMORY_BAR remains (an async load's DONE is its data arrival, not issue).

Model compute-output spad footprint in the SRAM version/capacity machinery so
buffer reuse (WAR) is capacity-modeled, not a hard edge. The output size comes
from the DMA records that touch the same buffer (a buf_bytes pre-pass); an
in-place buffer (accumulator, relu) is version-transparent so footprint is not
double-counted. The occupy gate and version release sit in the MOVIN/MOVOUT/COMP
issue points (release before the COMP skip path so a skipped matmul still frees).

Runtime: collapse child_inst / _pipeline_children into one event-indexed
_deps[ISSUE|DONE] with add_dep(c, on) and fire(e); collapse the weight-slot
release queue and the async-load wakeup into one _due_events timed-effect table
drained by process_due_events. Both are behavior-preserving (byte-identical).

Require the weight-slot model: sa_weight_buffer_depth must be > 0 (errors at
init), and the round-robin disable mode is removed. Degenerate traces (a
consumer-less preload, an unpinned matmul) hit explicit error+exit guards rather
than asserts that vanish under NDEBUG.

Mark the legacy ONNX TOG path deprecated: it is superseded by the trace path, so
TileGraphParser logs a deprecation warning and the TORCHSIM_LEGACY_TOG=1 opt-in
warns at command build.
… spad/2

The validation-binary spad-overflow check sat inside `if functional_mode:`, so in
timing-only / autotune (non-functional) runs an over-spad tiling was never
rejected and reached TOGSim, which wedged ("spad too small") and crashed the
compile via assert 0. Move the compile + check out of the functional gate (the
Spike execution itself stays gated, run_spike below) and budget per dispatch at
spad/2 -- two work-items run concurrently (double buffer), so each must fit half
the spad or they deadlock competing for it. This matches the GEMM tiling gate
(max_spad_size = spad/2), which pointwise ops lacked. Fixes the resnet /
test_scheduler wedge where a fused BatchNorm+ReLU tile exceeded spad/2.
Every generated trace.cpp now opens with a "GENERATED ... DO NOT EDIT" banner
carrying the TOGSim ABI version and the togsim_* call formats, so a dumped trace
is self-documenting. Both are read from togsim_runtime.h at codegen time -- the
version from the TOGSIM_ABI_VERSION define, the call-format text from a marked
block next to the declarations -- so they never drift when the ABI changes.
…rint

The bridge sums each work-item's distinct-buffer footprint (each buffer once, so
a reduction's reloads of the same section do not inflate it) onto its Tile.
can_issue then admits two concurrent dispatches only when each fits half the
spad; a dispatch larger than spad/2 runs alone with the whole spad, so two
work-items never compete for the shared spad and deadlock -- a runtime safety net
beneath the codegen spad/2 gate. The footprint and resulting max_dispatch are
logged on the TILE_SCHEDULED trace line for debugging.
…election

select_tile fed the epilogue buffer count to gemm_combination_mapping as
max(n_extra_read - 2, 0), which optimistically assumed the X and W DMA buffers
were already freed and reusable by the epilogue. The codegen lays every buffer
out as a disjoint .spad global and never reuses freed space, so the estimate
undercounted the real footprint: for a fused matmul+relu kernel the budget came
to 64512 B/lane (treated as a bare GEMM) while the emitted kernel used 89088
B/lane including the relu output buffer. Under the full-spad guard this was
harmless (89088 < 131072), but the spad/2 guard rejected it and crashed
test_transformer_fusion with SpadOverflowError, since the heuristic path has no
tile-shrink fallback.

Pass n_extra_node + n_extra_read instead: one output-tile-sized buffer per
epilogue node plus one per extra read operand, matching what the codegen emits.
For the matmul+relu kernel the budget now equals the actual footprint exactly,
and tile selection picks TILE_M=128 (62976 B/lane) which fits spad/2.

Liveness-based spad reuse and broadcast over-allocation are tracked separately
as optimizations in issue #275.
The spad-overflow check summed the kernel stack frame into the per-lane
scratchpad usage (spad_usage = stack_size + spad_size) and rejected the
tiling when that exceeded the spad/2 budget. But only the .spad section
actually lives in the scratchpad -- it is pinned there by the
--section-start=.spad link option. The kernel stack is in main memory
(sp is set up by pk in the -m region, not at the scratchpad vaddr), so it
does not consume the scratchpad and must not be charged against it. The
scalar frame is also shared across lanes, not per-lane, so adding it
double-counted.

On small configs (8x8) this falsely rejected feasible tilings: the
wrapper3 conv2d/resnet/mistral kernels fit the 32 KB spad with room to
spare but were tripped over the spad/2 gate purely by the ~200-800 B
stack term, crashing compile with SpadOverflowError. Drop the stack
term; the .spad-only check still correctly rejects real buffer overflows
(e.g. sparsity, which is fixed separately by f05ac8a).
Spike v1.0.3 zero-inits the MVIN/MVOUT DMA address buffer so ROUNDUP
padding entries are skipped instead of dereferencing uninitialized
garbage, fixing the host store segfault on 8-lane (wrapper3) configs
where a tile's split axis is not a multiple of vlane_stride*n_vu.
Add pytorchsim_functional_verify_per_kernel, a sub-option of
pytorchsim_functional_mode that localizes the first compiled kernel whose Spike
output diverges from a CPU reference.

When enabled, the generated wrapper compares every realized buffer (the output of
each fused kernel) against a CPU "golden" computed once per call by running the
original aten graph (V.graph.module) on CPU with the same inputs, via an
fx.Interpreter that records each node's output by name. A buffer is mapped to its
originating fx node through V.graph.get_buffer().origin_node, so the check reports
the kernel, the originating op, the offending indices and the max abs diff, then
raises at the first divergence (stop-at-first). Comparison granularity is the
fused-cluster output, the finest observable in a fused pipeline.

Auto-disabled when functional mode is off (no Spike values to verify); the config
accessor AND-gates the key with functional mode. Codegen bakes the
verify_init/verify_check calls into the wrapper only when enabled at compile time,
so clear the codegen cache when toggling. Tolerances via
TORCHSIM_FUNCTIONAL_VERIFY_RTOL / _ATOL (default 1e-4).

extension_functional_verify.py holds the graph registry and the runtime
golden/compare logic; mlir_codegen_backend injects the calls at the wrapper level;
extension_config reads the new YAML key.
The BMM tile selector fed only n_extra_node to gemm_combination_mapping and
dropped the prologue's extra-read operands entirely, so a softmax-fused
attention matmul (value^T @ softmax(scores)) was tiled as a bare BMM. The
codegen lays the softmax max/sum operands out as their own disjoint
weight-tile-sized .spad globals (buf3/buf4) and never reuses freed space, so
the estimate undercounted the real footprint: on the 32x32 wrapper2 config
(16 KB/lane spad/2 budget) tile selection picked TILE_N=512, whose emitted
kernel used 16896 B/lane and overflowed the budget by 512 B, crashing
test_transformer with SpadOverflowError. Under the 128x128 full-budget config
the slack hid it. This mirrors the epilogue n_extra_read gap fixed for the
GEMM template (commit f05ac8a), now on the BMM prologue path.

Add an n_prologue_extra_read knob to gemm_combination_mapping that charges
each extra prologue-read operand as one weight-tile-sized (TILE_K x TILE_N)
.spad buffer, matching what the codegen emits, and have the BMM template count
those operands the same way codegen does (the one numel-matching main input
reuses the matmul-operand buffer; every other read gets its own global). Tile
selection now rejects TILE_N=512 (16896 B/lane) and picks a fitting tile
(8704 B/lane), so wrapper2 test_transformer passes the full Spike + allclose
check. The new parameter defaults to 0, leaving the GEMM and conv callers
unchanged.
Add docs/per-kernel-functional-verify.md (usage, options, mechanism,
limitations, code map) and point to it from CLAUDE.md: a one-line entry in the
debugging section and in the YAML knobs list.
Keep only a one-line note on what TOGSIM_ABI_VERSION guards; the per-bump
v1..v12 history was noise nobody reads.
[Frontend] Budget fused-prologue spad buffers in BMM tile selection
…s off

format_dma_inst_issued_trace_line and format_instruction_detail_line are only
ever used as the message argument to trace_instruction_line (a spdlog::trace
sink). Because the argument is evaluated eagerly at the call site, the kernel
builds a formatted std::string for every issued/finished instruction and then
spdlog drops it whenever the level is above trace (the default). Bail out of
both formatters when trace logging is disabled so the fmt::format work is
skipped entirely.

This removes wasted work but is not the simulation-speed bottleneck (that is the
per-cycle ready-queue rescan in Core::cycle, addressed separately); the conv
kernel wall time is unchanged within noise.
Core::cycle() walked every tile's ready-instruction list on every simulated
cycle to find one instruction to issue. Profiling an 8x8 conv kernel showed this
scan is ~96% of simulation time and, within it, the walk is >99%: the ready list
holds ~2000 instructions (mostly blocked on the SRAM-capacity / weight-slot
throttle), only one issues per cycle, and the blocked ones are never removed, so
the same ~2000 are re-examined every cycle (~2.6k list iterations/cycle, ~1e9
over 400k cycles). At small arrays this dominates because tiny tiles inflate both
the ready-list length and the number of DMA-wait stall cycles.

Gate the scan behind a per-Core _issue_dirty flag. The scan's outcome can only
change when the ready set grows or a resource frees, so set the flag on exactly
those events:
 - ready set grows: Instruction::dec_ready_counter, on enqueueing a now-ready
   instruction, sets the OWNING core's _issue_dirty via _owner_dirty_ref (wired
   in Core::issue when the tile is dispatched) -- per-core so a remote core's
   enqueue does not force every core to rescan;
 - SRAM freed (release_sram) and weight slot freed (apply_due) set _issue_dirty;
 - a new dispatch (issue) and a successful issue keep it set.
On a cycle with none of these the scan would re-walk the same blocked
instructions and issue nothing, so it is skipped. Instructions are never moved
between queues, so there is no re-admission churn.

Issue-identical: it never changes which instruction issues or when, only whether
the (side-effect-free under those conditions) scan runs. The 8x8 conv kernel
produces the same 403026 cycles as before, with its TOGSim wall dropping from
214.5s to 41.9s (~5x, 71% of scans skipped). A forced-scan invariant check (no
skipped cycle ever issues or inline-finishes a zero-cycle COMP) found 0
violations across 39 kernels spanning GEMM, BMM, softmax and conv. DMA-bound
kernels (the ones that hit the 6h CI cap) stall more and gain more.
In Core::cycle()'s issue scan, a COMP with compute_cycle == 0 finishes inline and
calls instructions.erase(it), but does not set `issued` or break -- so control
falls through to the `it++` at the end of the loop body, incrementing the
std::list iterator that erase() just invalidated. That reads a freed node and is
undefined behavior; it usually limps along because the freed node's next pointer
is briefly intact, but under a different allocator / sanitizer / build it can
crash or skip-or-duplicate the remaining ready instructions in that tile,
corrupting issue order.

Use the iterator erase() returns and continue, the standard erase-in-loop idiom.
No behavior change on the path that happened to work; it just removes the UB.
YWHyuk added 17 commits July 7, 2026 16:09
…X TOG

The main compile/sim path no longer generates or selects the legacy ONNX
Tile-Operation-Graph. extension_codecache emits only trace.so + trace_cycles.tsv
(the build-skip now keys on trace.so), and TOGSimulator.run_standalone always
drives TOGSim with --trace_so. The TORCHSIM_LEGACY_TOG opt-in is removed from the
frontend. The ONNX --models_list branch is kept solely for the STONNE sparse path
(extension_op.py); TOGSim's C++ ONNX parser is untouched (separate PR).

origins (which FX nodes a kernel came from) is preserved: logged per kernel run
and recorded as a trailing "# origins:" line in trace_cycles.tsv -- the legacy
ONNX TOG carried this as node metadata, and the C++ cycle-table loader stops at
the comment so the current parser is unaffected.

Also drop the dead tog_file param from mlir_gem5_compile_command, migrate
scripts/chiplet.sh to --trace_so/--cycle_table (the trace path stubs per-tensor
addresses and --attributes_list is no longer a Simulator option), and refresh
the CLAUDE.md TOG-generation notes.
Rework per-kernel codegen so the loop body is an ordered list of
load->compute->store steps (class Step + push_step + codegen_loops
iteration) instead of the fixed dma_loads/compute/dma_stores buffers.
A step may be compute-only or transfer-only; the compute_idx loop is
emitted only for steps that actually have compute content.

Drop the ad-hoc self.masks buffer: the reduction-tail mask (get_mask)
now emits into the compute buffer, since a mask is just compute.

Use the step model to express indirect (gather/scatter) access cleanly.
When the indirect index is produced in the same pass,
convert_indirect_indexing pushes a new step for the offset build, so the
index load, the offset build and the gather become separate steps bridged
through spad. This replaces the compute_dependecy -> dma_stores hack in
both convert_indirect_indexing and load(). push_step clones the CSE so the
name counter is shared but the dedup cache is per-step (each step is its
own compute_idx region). Validated: x[idx+2]+1 emits two steps and matches
CPU; add/matmul/reduce/softmax/layernorm and gather/scatter unchanged.
…ptor

Replace the affine.apply{indirect_access} symbol smuggle with an explicit
offset descriptor. convert_indirect_indexing returns the offset spad instead
of folding sympy.Symbol(out) into the index; emit_transfer carries it as a
togsim.transfer operand; decompose_transfer lifts that operand to a
memref.dma_start {indirect_offset = @spad_symbol} attribute (memref.dma_start
is a registered op and rejects an extra operand, but accepts the attribute);
lower_dma_to_gemmini reads the attribute and resolves the global for CONFIG4
(drops _find_indirect); build_skeleton adds the offset spad to the gather DMA
read_bufs so the offset-build -> gather dependency edge forms in the trace.

The index stays clean (base only). Validated on both paths: Spike functional
(computed-index gather + scatter allclose, pointwise/reduce regression 0) and
the C++ trace timing path end-to-end (allclose; gather togsim_dma read_bufs now
includes the offset spad).

Indirect addressing (scattered-DMA timing) in the new trace path is a separate
gap tracked in issue #284.
…ride, compute-loop multi-dim offset

Three refinements on top of the explicit offset descriptor:

1. Detect indirect access by an explicit symbol set instead of an "tmp"
   substring match. indirect_indexing now records str(index_var) in
   self.indirect_symbols; a _has_indirect(expr) helper tests the index
   free_symbols against it; the former "tmp"-string sites (store,
   get_dma_info, convert entry/indirect_dims/stride) use it. Removes the
   now-dead _find_indirect from lower_dma_to_gemmini.

2. Single indirect dim: pass the raw index spad and let the MVIN apply the
   gather stride per position (CONFIG4 offset_stride) instead of a
   vector_load/muli/vector_store round-trip that pre-multiplied the stride.
   emit_transfer carries offset_stride; decompose copies the attribute;
   lower programs CONFIG4 with it (was hardcoded to 1).

3. Multi indirect dim (e.g. x[ix, iy]): the offset is a sum of strided
   indices, which a single CONFIG4 channel cannot do, so the sum stays in
   the kernel -- but build it in the compute loop (chunked by
   compute_vec_size, not a single tile-wide vector) and store it to a
   DEDICATED offset spad so an index that is live elsewhere (x[ix, iy] + ix)
   is not clobbered. push_step separates the offset-build loop from the
   gather that reads it.

Adds multi-dim gather and index-reuse cases to test_indirect_access.
Validated: indirect/scatter/embedding + the two new multi-dim cases pass;
add/matmul/softmax regression 0.
…emmini

memref.dma_start is a registered op with a fixed operand list, so it cannot carry the
runtime descriptor operands the masked-DMA clamp needs (per-dim low/high are dynamic-shape
values, not attributes). togsim.transfer is unregistered, so it carries them. Eliminate
the dma_start intermediate: togsim.transfer flows through the whole pipeline and lowers
directly to Gemmini.

- New lower_transfer_to_gemmini merges decompose_transfer's <=4D handling (unit-dim
  collapse / >4D affine.for peel + lane-banked SRAM offset) with lower_dma_to_gemmini's
  CONFIG/CONFIG2/CONFIG3/[CONFIG4]/MVIN|MVOUT asm. A/B-validated (func7 + packed CONFIG
  rs1/rs2) against the old decompose->lower chain on add/pad/matmul/gather.
- togsim.transfer gains a tag_idx operand (dram,dram_idx,sram,sram_idx,tag,tag_idx,
  dma_type,vst[,offset]) so the async DMA<->barrier runtime slot survives without
  memref.dma_start's variadic tag indices.
- memref.dma_wait -> togsim.wait (togsim-level counterpart, emitted by the vcix matmul
  lowering, erased by the merged lowering). No memref.* DMA ops remain.
- decompose_transfer dropped from PRE_OPT; loop-padding runs OPAQUELY on togsim.transfer
  (mlir-opt --allow-unregistered-dialect), doing only its loop-bound rounding (its
  DRAM-grow half is what the future masked-DMA clamp replaces).
- dma_fine_grained, lower_to_vcix._DmaView, build_skeleton ported to read togsim.transfer.

Functional path (Spike, timing_mode=0) validated: add, matmul (fine-grained subtile +
async + togsim.wait), softmax, gather/scatter (indirect CONFIG4), constant_pad, conv
(padding=0) all allclose. Known-pending: conv with padding>0 regresses -- it relied on
loop-padding's DRAM-grow, which the masked-DMA clamp (next step) replaces (padding=0
passes; wrong values, not a crash).
…wait

Finishes the memref.dma_start removal for the trace/timing path. build_tog's
TogBuilder now decodes togsim.transfer (dram,dram_idx,sram,sram_idx,tag,tag_idx,
dma_type,vst[,offset]) and togsim.wait (tag,tag_idx,num_elements) instead of
memref.dma_start/dma_wait: op dispatch, _dma_start_fields (positional decode +
direction from dma_kind / dma_type), the tag-user base-address scan, and the wait
operand reads. TOGDMANode.tile_stride keeps reading the dram_stride attr (unchanged
semantics). build_skeleton's wait marker/dispatch move to togsim.wait so barriers
still become togsim.memory_barrier. dep_analysis needs no change (the decoder keeps
the same field keys).

Trace/timing (TOGSim, timing_mode=1) validated: add 5/5, matmul 11/11 (async DMA +
tag_idx barrier pairing). No memref.* DMA ops remain in either path.
… not operand count

len(operands) > 8 is ambiguous once the masked-DMA clamp adds more optional operands
(low/high after offset). Gate on the `indirect` attr instead, so offset vs mask operands
are disambiguated by explicit flags. Behavior-preserving: gather/add pass on both the
functional (Spike) and trace (TOGSim) paths.
…gsim.wait)

The trace producer now reads togsim.wait (not memref.dma_wait); update the pass
docstring and comments to match. No code change.
…o CONFIG/2/3/4)

Instead of packing the transfer params into four CONFIG asm instructions, build a
144-byte DMA descriptor as an 18xi64 memref.global (.data, non-constant so runtime
fields can be written) and emit a single CONFIG_DESC that hands its address to the
MVIN/MVOUT (which read the struct, matching the Spike-side torchsim_config_desc). This
is the vehicle for the masked-DMA low/high clamp: dim_low/dim_high/fill are descriptor
fields (currently defaulted, so behavior is unchanged), not new special registers.

- _pack_desc encodes the fields into the struct's i64 slots (dim_size/low/high/mm_stride/
  spad_stride/element_size/vlane/config_type/flags/indirect stride+esize).
- descriptor globals are deduped by content; the indirect offset spad address is a
  runtime value, stored into slot 15 before the transfer (flags.bit0 = indirect).

Requires the Spike torchsim_config_desc instruction (riscv-isa-sim). Validated
functional (Spike) + trace (TOGSim): add, matmul (subtile+async), softmax,
gather/scatter (indirect runtime store), constant_pad all pass.
The togsim.transfer descriptor lowering requires the spike torchsim_config_desc
instruction (riscv-isa-sim PR #5, tag v1.0.4). Track the new tag so build-from-source
picks up the descriptor-capable spike.
The masked-DMA clamp needs spike to read the descriptor dim_low/dim_high +
desc_fill, and index_add needs the MVOUT accumulate flag. Pin v1.0.5.
…-dividing tails

Drop the index_expr / indirect tile-divisibility recompile so pad_vlane_tile's
lane-aligned tiles are kept instead of shrunk to a divisor of the dim. A
non-dividing tile's per-tile remainder is handled by a masked-DMA clamp:

- _emit_clamp emits per tile axis a dynamic [low, high) as affine.min/max of that
  axis' loop iv (the tile base), so the last partial tile and the pad borders fall
  out per iteration; low = max(0, glo - base), high = min(tile, ghi - base). The
  pointwise path (_masked_bounds) derives [glo, ghi) per dim (MVOUT [0, output);
  MVIN of a padded input [pad, pad+input_extent)); the MLIR templates (def_dma_op,
  e.g. gemm) derive it from the DRAM extent so a matmul tile padding past M/N/K is
  filled with 0 instead of MAC-ing garbage.
- emit_transfer carries the (low, high) index operands + masked_axes;
  lower_transfer_to_gemmini writes the runtime low/high into the descriptor
  dim_low/dim_high through an i32-view (byte layout unchanged) + masked flag +
  desc_fill. One desc_ptr helper builds every descriptor: a fully static one is a
  deduped i64 global, while any runtime-written field (masked dim_low/dim_high
  and/or the indirect address, stored as two i32 words) uses a unique i32-view
  global, so a single descriptor carries both a masked clamp and a gather/scatter,
  and indirect DMAs drop their divisibility recompile too.
- The box-excluded fill is the consuming reduction identity (reduction_init: sum 0,
  prod 1, max -inf, min +inf; 0 otherwise -- matmul MAC and stores use 0). Log-sum-
  exp exception (softmax / log_softmax): a sum whose node came from exp reduces
  exp(input), so the PRIMARY input (indexed by a reduction itervar) is filled with
  -inf; broadcast operands (the subtracted max) keep their finite identity, else
  input - (-inf) = +inf -> exp = inf -> NaN.
- index_add: instead of the compute gather-add-overwrite (which loses duplicate
  indices landing in the same non-dividing tile), the MVOUT accumulates out[idx] +=
  val via an accumulate flag; the DMA processes positions sequentially so duplicates
  accumulate correctly regardless of tile size. acc_float selects float vs integer
  add (element_size alone is ambiguous).

With no caller left, is_dim_dividable / adjust_tile_to_divisible and the
TileConstraint.must_divide_dim machinery (trim_large_tail bias, adjust loop) are
removed.

Fixes elementwise, constant_pad (mobilenet ragged vlane holes), reductions
(sum/max/min/prod/softmax/log_softmax/var/std/layernorm/mean), gather/scatter
(index_add with duplicates) and matmul/addmm on non-dividing shapes. Requires spike
with the descriptor dim_low/dim_high clamp, desc_fill and the MVOUT accumulate flag.
…ne-grained)

Conv on a non-dividing reduction (K = in_ch * kh * kw, e.g. 3*7*7=147 > systolic
size) MAC-ed garbage from the tail of the K tile. Two fixes:

- def_dma_op now takes the per-dim clamp extent from the kernel's loop_extents
  {loop_iv_name: loop_bound}, which each conv template declares once (the template's
  ranges/itervars equivalent). The stride-match extent inference is unreliable for
  conv because the gemm reads a REPACKED weight (O_C contiguous, stride 1) whose
  strides do not match the declared node_layout, so the inferred extent is wrong
  (e.g. O_C=64 read as 8); gemm (no loop_extents) keeps the inference. Only a dim
  whose index is a single iv is clamped -- an im2col spatial axis (o_h + k_h) has no
  single base and is left unclamped, which is correct because that input is
  pre-padded. All four conv templates (base, mt, sb, sbs) declare loop_extents; note
  the mt template fuses channel and kernel-width so its tile_k loops 0..I_C*K_W.
- dma_fine_grained now PROPAGATES the masked (and indirect / accumulate) attrs +
  the low/high operands when it splits a tile into subtiles, and REMAPS each clamp
  per subtile: a subtile at offset iv*sub sees high = min(sub, high - iv*sub),
  low = max(0, low - iv*sub) (out-of-window subtiles get a negative high / low > sub
  and Spike skips them). Previously it dropped the clamp, so the conv clamp never
  reached Spike; gemm was unaffected because its DMAs do not go through this pass.

Fixes conv2d on non-dividing in_channels / kernels (test_conv2d 13/13), including
fused pointwise epilogues (conv + relu/scale/bias/sigmoid). No regression:
matmul/addmm, bmm, elementwise, reduce, softmax, index_add all pass.
A fused template epilogue (matmul/conv + relu/add/gelu/...) stores its result via
store_epilogue, which -- unlike the plain store() -- emitted no masked-DMA clamp. On
a non-dividing output the store then wrote the systolic pad region past the real
extent, corrupting the result (fused matmul+relu at 10x10x10 was off by ~7).

store_epilogue now clamps each output tile dim to [0, extent): the extent of every
loop iv is ranges[k], and the tile dims are indexed by dim_aliasing (tile-dim order),
so a dim whose iv extent does not divide the tile is clamped. Keyed by iv name, it is
naming-agnostic -- it covers both the index-N matmul epilogue and the c0/tile_n/...
conv epilogue -- and _emit_clamp skips dividing dims, so a dividing output is a no-op.

Fixes fused matmul/conv epilogues on non-dividing shapes (matmul+relu/add/gelu,
conv+relu/scale/bias/sigmoid, non-dividing K and non-dividing output). No regression:
test_matmul, matmul_activation, addmm_residual, matmul_scalar all pass.
Hand-write a masked MVOUT togsim.transfer, run lower_transfer_to_gemmini in-process
(no torch, no spike; skipped without the MLIR bindings), and compare to an embedded
golden -- the exact lowered IR (MLIR SSA numbering is deterministic for a fixed
input): the i32-view descriptor global (dim_size / config_type / masked flag
packed), the runtime dim_low/dim_high stores at i32 idx 6/10 for masked_axes=[2],
and the CONFIG_DESC + MVOUT asm. A second case adds subtile_size=[1,1,4,8] and shows
the descriptor dim_size become the subtile (config) shape. Timing mode erases the
transfer. Under tests/lowering/ so further per-pass goldens have a home. Regenerate
the golden and review the diff on an intended lowering change.
…dims, cleanup

Addresses the PR review of the masked-DMA non-dividing work:

- _masked_fill_bits detected a log-sum-exp reduction with `"exp" in origin.name`, which
  also matches expand / expm1 / experimental -- an ordinary non-dividing sum whose origin
  name contains "exp" would then get a -inf tail fill instead of the sum identity 0 and
  reduce to -inf. Match the op target instead (_origin_is_exp: the target's overload
  packet name == "exp"); same fix applied to the get_padding_type twin. Unit test in
  tests/lowering/test_masked_fill.py pins that expand/expm1 do not match.
- Remove the self._last_local_dims temporal coupling: get_dma_info now returns local_dims
  and load()/store() pass it to _masked_bounds explicitly, so a reordered or skipped
  get_dma_info can no longer feed stale clamp axes.
- Drop the underscore-prefixed locals in def_dma_op / store_epilogue (non-idiomatic),
  trim the oversized _masked_fill_bits / _masked_bounds docstrings to the invariant plus a
  link to docs/design/masked-dma-descriptor.md, and clarify the bf16 dtype-table comment
  (present only for table totality; bf16 is not actually supported).
Dropping the tile-divisibility recompile means non-dividing shapes rely entirely on the
masked-DMA clamp/fill; a gap is now a silent wrong answer, not a slow-but-safe recompile.
Add tests/ops/misc/test_masked_nondividing.py pinning that path: 2-D constant pad (per-dim
greedy pad recovery), sum over a broadcast operand (0 fill, guards the exp-substring bug),
softmax (genuine -inf fill), non-dividing elementwise, and non-dividing gather. Register it
in the CI allowlist (pytorchsim_test.yml). Also add the missing trailing newline to
test_indirect_access.py.
@YWHyuk YWHyuk force-pushed the feature/togsim-cpp-trace branch from 1e43e9a to cfcb7a9 Compare July 7, 2026 07:10
@YWHyuk YWHyuk force-pushed the feature/togsim-dynamic-shape branch 2 times, most recently from 9dbc257 to 908ea30 Compare July 7, 2026 07:37
YWHyuk added 4 commits July 7, 2026 19:37
…asked-DMA clamp

Make the MLIR backend emit valid IR under torch.compile(dynamic=True), where a
loop range is a sympy symbol (e.g. s52) instead of a concrete int.

- is_symbolic_dim: single predicate for "runtime dim"; the tile-fit heuristics
  (trim_large_tail, get_padding_ratio, make_choices) skip a symbolic dim rather
  than doing concrete-int arithmetic on it (kept fixed init tile; tail becomes a
  runtime remainder). get_mlir_shape emits a dynamic memref dim (memref<?>).
- LoopLevel emits a symbolic upper bound as an index SSA (%<name>_bound);
  codegen_loops materializes it (memref.load + index_cast) in the prologue.
- set_ranges registers each symbolic dim's runtime-extent SSA once
  (dyn_bound_cses), the single source of %<name>_bound (dyn_bound_name); the
  loop bound and the masked-DMA clamp both read that var instead of rebuilding
  the name at codegen time.
- Masked-DMA clamp: a symbolic tile axis emits the same per-dim [low, high)
  clamp as a static non-dividing axis, with the runtime extent as the affine.min
  symbol operand -- so a non-dividing dynamic shape is handled by one path and
  the last partial tile falls out at runtime.
- axis_split: aligned axis-split detection is symbolic-aware (+ unit test).

All guards gate on is_symbolic_dim, so the static path is unchanged.
…gnostic validation, producer shape_args

Drive a dynamic kernel through compile + timing + functional validation without
a concrete extent.

- cycle_table.pin_loops_to_one_tile: per-tile cost is shape-invariant, so sample
  cycles on a one-tile copy of the (symbolic) IR; extension_codecache sizes the
  sampling host buffers to one tile (_concretize_attrs_for_sampling).
- mlir_caller_codegen: the Spike validation binary is shape-agnostic -- it loads
  each size arg into N_<sym> and mallocs / builds the memref descriptors from N
  at runtime, so one binary serves any size. extension_codecache builds and runs
  it for a dynamic kernel too; the per-tile spad measurement stays valid, so the
  spad-overflow gate runs unconditionally (static path unchanged).
- lower_to_emitc: the trace producer re-sources its loop bound from shape_args,
  so one trace.so serves any size. build_tog/build_skeleton skip the bfs
  serialization that assumes a concrete loop_end.
- is_dynamic_shape is detected from the size-symbol arg (MLIR_ARGS_VAR), not an
  IR-text sniff.
…via the attribute file

The dynamic trace producer reads its loop bounds from shape_args; feed them at
simulation time through the existing per-kernel attribute YAML (which already
carries address_info), not a bespoke channel.

- simulator: a scalar input (a dynamic size arg, e.g. s52) is not a tensor
  address -- collect such scalars into a shape_args sequence in the YAML, in arg
  order (== the producer's shape_args[k]); pass --attribute alongside --trace_so.
  dump_args/write_arg write the size arg's runtime value (int64) so the
  shape-agnostic Spike binary can load its loop bound.
- main.cc: add --attribute; in the trace branch load the YAML and fill
  shape_args, passed to run_producer (was nullptr, 0).
A single torch.compile(dynamic=True) add compiles to one trace producer .so and
is simulated at several input sizes -- the producer reads its loop bound from
shape_args at runtime, so the same .so serves any size. Exercises the whole
dynamic-shape pipeline end to end, including non-dividing sizes (the masked-DMA
symbolic clamp handles the remainder tile).
@YWHyuk YWHyuk force-pushed the feature/togsim-dynamic-shape branch from 908ea30 to 0797c45 Compare July 7, 2026 10:37
@YWHyuk YWHyuk force-pushed the feature/togsim-cpp-trace branch 3 times, most recently from d9131ba to 1ff2ebb Compare July 10, 2026 13:00
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant